from typing import Any, Type, Union

import numpy as np
import torch
import torch.nn as nn
from gym.spaces import Box, Discrete
from torch.distributions import Categorical, Normal
from typing_extensions import override

from ....network.utils import make_mlp


class Actor(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def _distribution(self, obs) -> Union[Normal, Categorical]:
        raise NotImplementedError

    def _log_prob_from_distribution(self, pi, act):
        raise NotImplementedError

    def forward(self, obs, act=None):
        pi = self._distribution(obs)
        logp_a = None
        if act is not None:
            logp_a = self._log_prob_from_distribution(pi, act)
        return pi, logp_a

    def distribution_param(self, obs):
        raise NotImplementedError


class MLPCategoricalActor(Actor):
    def __init__(self, **kwargs):
        super().__init__()
        self.logits_net = make_mlp(
            [kwargs["obs_dim"]] + list(kwargs["hidden_sizes"]) + [kwargs["act_dim"]],
            kwargs["activation"],
        )

    def _distribution(self, obs):
        logits = self.logits_net(obs)
        return Categorical(logits=logits)

    def _log_prob_from_distribution(self, pi, act):
        return pi.log_prob(act)

    def distribution_param(self, obs):
        return self.logits_net(obs)


class MLPGaussianActor(Actor):
    def __init__(self, **kwargs):
        super().__init__()
        self.log_std = torch.nn.Parameter(
            -0.5 * torch.ones(kwargs["act_dim"], dtype=torch.float32)
        )
        self.mu_net = make_mlp(
            [kwargs["obs_dim"]] + list(kwargs["hidden_sizes"]) + [kwargs["act_dim"]],
            kwargs["activation"],
        )

    def _distribution(self, obs):
        mu = self.mu_net(obs)
        std = torch.exp(self.log_std)
        return Normal(mu, std)

    def _log_prob_from_distribution(self, pi, act):
        return pi.log_prob(act).sum(axis=-1)

    def distribution_param(self, obs):
        return self.mu_net(obs)


class Critic(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, obs):
        ...


class MLPCritic(Critic):
    def __init__(self, **kwargs):
        super().__init__()
        self.v_net = make_mlp(
            [kwargs["obs_dim"]] + list(kwargs["hidden_sizes"]) + [1],
            kwargs["activation"],
        )

    def forward(self, obs):
        return torch.squeeze(self.v_net(obs), -1)


class ActorCritic(nn.Module):
    def __init__(self, Actor: Type[Actor], Critic: Type[Critic], **kw_args) -> None:
        super().__init__()
        self.pi = Actor(**kw_args)
        self.v = Critic(**kw_args)

    def process_obs(self, obs) -> Any:
        raise NotImplementedError

    def process_obs_batch(self, obs) -> Any:
        raise NotImplementedError

    def process_act(self, act) -> Any:
        raise NotImplementedError

    def process_act_batch(self, act) -> Any:
        raise NotImplementedError

    def step(self, obs, random: bool):
        self.eval()
        with torch.no_grad():
            obs = self.process_obs(obs)
            pi = self.pi._distribution(obs)
            a = pi.sample() if random else pi.mean
            a = self.process_act(a)
            logp_a = self.pi._log_prob_from_distribution(pi, a)
            v = self.v(obs)
        return a.cpu().numpy(), v.cpu().numpy(), logp_a.cpu().numpy()

    def act(self, obs):
        return self.step(obs, True)[0]


class MLPActorCritic(ActorCritic):
    def __init__(
        self, observation_space, action_space, hidden_sizes=(64, 64), activation=nn.Tanh
    ):
        obs_dim = observation_space.shape[0]

        if isinstance(action_space, Box):
            Actor = MLPGaussianActor
            act_dim = action_space.shape[0]

        elif isinstance(action_space, Discrete):
            Actor = MLPCategoricalActor
            act_dim = action_space.n

        Critic = MLPCritic

        super().__init__(
            Actor,
            Critic,
            obs_dim=obs_dim,
            act_dim=act_dim,
            hidden_sizes=hidden_sizes,
            activation=activation,
        )

    def process_obs(self, obs):
        return torch.as_tensor(obs, dtype=torch.float32)

    def process_obs_batch(self, obs) -> Any:
        obs_batch = np.array([np.array(item) for item in obs])
        return torch.tensor(obs_batch, dtype=torch.float32)

    def process_act(self, act) -> Any:
        return torch.as_tensor(act, dtype=torch.float32)

    def process_act_batch(self, act) -> Any:
        act_batch = np.array([np.array(item) for item in act])
        return torch.tensor(act_batch, dtype=torch.float32)
